#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os,sys,re
import MySQLdb
from wwpy.util.Module import sqlite,postgre
from util.Logger import logger
from Query import Query
from DBC import DBC,dbc

class DBA(object):
    def __init__(self,coding='utf8'):
        self.coding=coding
        self.connection=None
        
    def set_coding(self,coding='utf8'):
        if coding is not None:
            self.coding=coding
        self.execute("SET NAMES %s" % self.coding)
        
    def connect(self):
        raise Exception('Not implement.')
    
    @property
    def is_connected(self):
        if self.connection is None:
            return False
        return True
    
    def execute(self,sql,*args,**kwargs):
        if not self.is_connected:
            self.connect()
        if hasattr(self.connection,'execute') and callable(self.connection.execute):
            self.cursor=self.connection.execute(sql,*args,**kwargs)
        elif hasattr(self.connection,'cursor') and callable(self.connection.cursor):
            self.cursor=self.connection.cursor()
            self.cursor.execute(sql,*args,**kwargs)
        elif callable(self.connection):
            self.cursor=self.connection()
            self.cursor.execute(sql,*args,**kwargs)
        else:
            raise Exception('Could not find callable execute method from connection or cursor!')
        if hasattr(self.connection,'commit') and callable(self.connection.commit):
            self.connection.commit()
        return self.cursor
    
    def do(self,sql,*args,**kwargs):
        try:
            self.execute(sql,*args,**kwargs)
        except Exception,error:
            raise Exception(error)
        
    def fields(self,sql=None):
        if sql is not None:
            self.execute(sql)
        if hasattr(self.cursor,'description'):
            fields=tuple(map(lambda a:a[0],self.cursor.description))
        elif hasattr(self.cursor,'keys') and callable(self.cursor.keys):
            fields=tuple(self.cursor.keys())
        else:
            raise Exception('Could not determine field names from cursor: %r' % self.cursor)
        return tuple(fields)
    
    def row(self):
        if hasattr(self.cursor,'fetchone') and callable(self.cursor.fetchone):
            return self.cursor.fetchone()
        return []
        
    def rows(self):
        if hasattr(self.cursor,'fetchall') and callable(self.cursor.fetchall):
            return self.cursor.fetchall()
        return []
    
    def query(self,sql,*args,**kwargs):
        rows=self.execute(sql,*args,**kwargs)
        fields=self.fields()
        for row in rows:
            yield dict(zip(fields,row))
            
    def records(self,sql,*args,**kwargs):
        return self.query(sql,*args,**kwargs)
    
    def results(self,has_fields=False):
        results=[]
        if has_fields:
            fields=self.fields()
            results.append(fields)
        rows=self.rows()
        results.extend(rows)
        return results
        
    def close(self):
        if hasattr(self.cursor,'close') and callable(self.cursor.close):
            self.cursor.close()
        if hasattr(self.connection,'close') and callable(self.connection.close):
            self.connection.close()
        self.connection=None
    
    def fetch(self,sql,has_fields=False,has_rows=True,row_func=None,*args,**kwargs):
        try:
            if not has_rows:
                return self.do(sql,*args,**kwargs)
            self.execute(sql,*args,**kwargs)
            results=self.results(has_fields)
            if row_func is not None:
                return self.memorize(results,row_func)
            return results
        except Exception,error:
            raise Exception(error)
    
    def row_func(self,info,row):
        info=info or []
        info.append(row)
        return info
    
    def memorize(self,results,row_func=None):
        info=None
        row_func=row_func or self.row_func
        for row in results:
            info=row_func(info,row)
        return info
    
    def scan(self,results,sep='\t'):
        for row in results:
            print sep.join(str(field) for field in row)
    
    def store(self,results,file,sep='\t'):
        f=open(file,'w')
        for row in results:
            f.write(sep.join(str(field) for field in row)+'\n')
        f.close()
        
    def add_conditions(self,conditions):
        if isinstance(conditions,(basestring,tuple,list,dict)):
            if isinstance(conditions,basestring):
                cdts=conditions
            elif isinstance(conditions,(tuple,list)):
                cdts=' AND '.join(conditions)
            else:
                def _set_kv(k,v):
                    if v is None:
                        return "%s IS NULL" % k
                    elif isinstance(v,(int,long,float)):
                        return "%s=%s" % (k,v)
                    else:
                        return "%s='%s'" % (k,v)
                cdts=' AND '.join(map(lambda (k,v):_set_kv(k,v),conditions.iteritems()))
            return cdts
        return None

class MySQL(DBA,Query):

    def __init__(self,host='localhost',port=3306,user=None,passwd=None,db=None,coding='utf8',**kwargs):
        self.host=host
        self.port=int(port)
        self.user=user
        self.passwd=passwd
        self.db=db
        Query.__init__(self)
        DBA.__init__(self,coding)
        
    def get_dbc(self):
        return dict(
            host=self.host,
            port=self.port,
            user=self.user,
            passwd=self.passwd,
            db=self.db,
            coding=self.coding,
        )
        
    def set_dbc(self,**dbc):
        self.host=dbc.get('host',self.host)
        self.port=int(dbc.get('port',self.port))
        self.user=dbc.get('user',self.user)
        self.passwd=dbc.get('passwd',self.passwd)
        self.db=dbc.get('db',self.db)
        self.coding=dbc.get('coding',self.coding)
        self.connection=None
        
    def connect(self):
        self.connection=MySQLdb.connect(host=self.host,port=self.port,user=self.user,passwd=self.passwd,db=self.db)
        self.set_coding(self.coding)
        
    def create(self,table,fields,is_update=False,options={'ENGINE':'InnoDB'},defaults={'CHARSET':'utf8'}):
        if self.has_table(table) and (not is_update):
            return
        sql="CREATE TABLE %s " % table
        if isinstance(fields,(tuple,list)):
            sql+='('+','.join(fields)+') '
        sql+=' '.join(map(lambda (k,v):'='.join([k,str(v)]),options.items()))
        sql+=' DEFAULT %s ' % ' '.join(map(lambda (k,v):'='.join([k,str(v)]),defaults.items()))
        self.do(sql)
        
    def load(self,file,table,fields=None,
             priority='LOW_PRIORITY',local='LOCAL',cover='IGNORE',
             field_sep='\\t',optionally=False,enclosed_char='',escaped_char='',
             line_sep='\\n',started_char='',ignore=None,others=None,sets=None):
        sql="LOAD DATA %s %s INFILE '%s' %s INTO TABLE %s " % (priority,local,file,cover,table)
        if isinstance(fields,(list,tuple)):
            sql+="(%s) " % ','.join(fields)
        sql+="FIELDS TERMINATED BY '%s' " % (field_sep)
        if optionally:
            sql+="OPTIONALLY "
        if enclosed_char:
            sql+="ENCLOSED BY '%s' " % enclosed_char
        if escaped_char:
            sql+="ESCAPED BY '%s' " % escaped_char
        sql+='LINES '
        if started_char:
            sql+="STARTING BY '%s' " % started_char
        sql+="TERMINATED BY '%s' " % line_sep
        if isinstance(ignore,int):
            sql+="IGNORE %s LINES " % ignore
        if isinstance(others,(list,tuple)):
            sql+="(%s) " % ','.join(others)
        if isinstance(sets,(list,tuple)):
            sql+="SET %s " % ','.join(sets)
        logger.info("Load file %s into table %s." % (file,table))
        self.do(sql)
        
    def save(self,file,table,table_fields=None,sep='\\t',is_drop=False,is_update=True,is_split=False,is_check=True,**kwargs):
        files=file if isinstance(file,(tuple,list)) else [file]
        options=kwargs.pop('options',{'ENGINE':'InnoDB'})
        defaults=kwargs.pop('defaults',{'CHARSET':'utf8'})
        if is_drop:
            self.delete(table,is_drop)
        if table_fields is not None:
            self.create(table,table_fields,options=options,defaults=defaults)
        where=kwargs.pop('where',None)
        where=self.add_conditions(where)
        if (not is_drop) and is_update:
            self.delete(table,conditions=where)
        if (not is_drop) and is_check:
            sql="SELECT COUNT(1) FROM %s WHERE 1=1 " % table
            if where:
                sql+='AND %s ' % where
            rows=self.fetch(sql)
            count=int(rows[0][0])
            if count>0:
                return
        part_lines=kwargs.pop('part_lines',1000000)
        if is_split:
            from wwpy.file.Spliter import Spliter
            part_files=[]
            for file in files:
                s=Spliter(file)
                parts=s.split(by='line',part_lines=part_lines)
                part_files.extend(parts)
            files=part_files
        from wwpy.file.File import File
        for file in files:
            f=File(file)
            if f.is_compressed:
                dest_file=f.decompress()
                self.load(dest_file,table,field_sep=sep,**kwargs)
                f.set(dest_file)
                f.remove()
            else:
                self.load(file,table,field_sep=sep,**kwargs)
        if is_split:
            map(os.remove,files)
            
    def dump(self,sql,file=None,has_fields=False,sep='\\t'):
        if file is None:
            file='dump_records_from_sql.txt'
        rows=self.fetch(sql,has_fields)
        f=open(file,'w')
        for row in rows:
            f.write(sep.join(str(field if field is not None else '') for field in row)+'\n')
        f.close()
        
    def delete(self,table,is_drop=False,conditions=None,order=None,limit=None):
        tables=table if isinstance(table,(list,tuple)) else [table]
        if is_drop:
            sql="DROP TABLE IF EXISTS %s " % ','.join(tables)
        else:
            sql="DELETE FROM %s WHERE 1=1 " % table
            if conditions is not None:
                sql+='AND %s ' % self.add_conditions(conditions)
                if limit is not None:
                    if order_by is not None:
                        if isinstance(order,(tuple,list)):
                            order=','.join(order)
                        sql+="ORDER BY %s " % order
                    sql+="LIMIT %s " % limit
        self.do(sql)
        
    def insert(self,table,fields=None,values=None,select=None,sets=None,others=None):
        sql="INSERT INTO %s " % (table,)
        if isinstance(fields,(tuple,list)):
            sql+='('+','.join(fields)+') '
        if isinstance(values,(tuple,list)):
            if isinstance(values,tuple):
                values=[values]
            def _set_value(v):
                if v is None:
                    return 'NULL'
                elif isinstance(v,(int,float,long)):
                    return str(v)
                return "'%s'" % str(v).replace("'","\\'")
            sql+="VALUES "+','.join("("+",".join(_set_value(v) for v in value)+")" for value in values)+" "
        elif isinstance(select,basestring):
            sql+="%s " % select
        elif isinstance(sets,dict):
            sql+="SET %s " % ','.join('='.join([k,str(v)]) for k,v in sets.items())
        if isinstance(others,basestring):
            sql+="%s " % others
        self.do(sql)
        
    def update(self,table,sets,conditions=None,order=None,limit=None):
        sql="UPDATE %s SET %s WHERE 1=1 " % (table,','.join('='.join([k,str(v)]) for k,v in sets.items()))
        conditions=self.add_conditions(conditions)
        if conditions:
            sql+='AND %s ' % conditions
        if limit is not None:
            if order is not None:
                if isinstance(order,(tuple,list)):
                    order=','.join(order)
                sql+="ORDER BY %s " % order
            sql+="LIMIT %s " % limit
        self.do(sql)
        
    def table_fields(self,table,cols=[0]):
        sql="SHOW FULL FIELDS FROM %s " % table
        rows=self.fetch(sql)
        if cols is None or not isinstance(cols,(tuple,list)):
            cols=[0]
        if cols==[0]:
            return tuple(row[0] for row in rows)
        return tuple(tuple(map(lambda col:row[col],cols)) for row in rows)
        
    def has_db(self,db):
        db=str(db).lower()
        sql="SHOW DATABASES"
        dbs=map(lambda t:str(t[0]).lower(),self.fetch(sql))
        if db in dbs:
            return True
        return False
        
    def has_table(self,table):
        table=str(table).lower()
        sql="SHOW TABLES"
        tables=map(lambda t:str(t[0]).lower(),self.fetch(sql))
        if table in tables:
            return True
        return False
    
    def has_record(self,table,where=None):
        if not self.has_table(table):
            return False
        sql="SELECT COUNT(1) FROM %s WHERE 1=1 " % (table)
        where=self.add_conditions(where)
        if where:
            sql+='AND %s ' % where
        records=map(lambda t:t[0],self.fetch(sql))
        if int(records[0]):
            return True
        return False
    
    def batch(self,sql,once_num=1000000,has_fields=False):
        sql=sql.strip(' ; ')
        regex=re.compile(r'(select) .+? (from .*)',re.I)
        m=regex.search(sql)
        if m:
            new_sql="%s COUNT(1) %s" % (m.group(1),m.group(2))
            rows=self.fetch(new_sql)
            record_num=int(rows[0][0])
        else:
            record_num=once_num
        if record_num<=once_num:
            rows=self.fetch(sql,has_fields)
            for row in rows:
                yield row
        else:
            print "SQL record number is %s." % record_num
            num=int(record_num/once_num)
            for n in range(num+1):
                _sql=sql
                offset=n * once_num
                _sql+=" LIMIT %s,%s" % (offset,once_num)
                print "Execute SQL: %s" % _sql
                rows=self.fetch(_sql,has_fields)
                for row in rows:
                    yield row
        
class SQLite(DBA):

    def __init__(self,db=None,coding='utf8',**kwargs):
        self.db=db
        super(SQLite,self).__init__(coding)
        
    def set_coding(self,coding='utf8'):
        if coding is not None:
            self.coding=coding
        self.execute('PRAGMA encoding="%s"' % self.coding)
        
    def connect(self):
        self.connection=sqlite.connect(database=self.db)
        self.set_coding(self.coding)
        
    def create(self,table,fields=None,select=None,temp=False):
        sql="CREATE "
        if temp:
            sql+="TEMP " 
        sql+="TABLE IF NOT EXISTS %s " % table
        if isinstance(fields,(tuple,list)):
            sql+="(%s) " % ','.join(fields)
        elif isinstance(select,basestring):
            sql+="AS %s " % select
        else:
            raise Exception('At least one param (fields or select) is not None.')
        self.do(sql)
        
    def insert(self,table,fields=None,values=None,select=None):
        sql="INSERT INTO %s " % table
        if isinstance(fields,(tuple,list)):
            sql+="(%s) " % ','.join(fields)
        if isinstance(values,(tuple,list)):
            sql+="VALUES ('%s') " % ("','".join(str(value).replace("'","\\'") for value in values))
        elif isinstance(select,basestring):
            sql+=select
        else:
            raise Exception('At least one param (fields or select) is not None.')
        self.do(sql)
        
    def delete(self,table,conditions=None,is_drop=False):
        if is_drop:
            sql="DROP TABLE IF EXISTS %s " % table
        else:
            sql="DELETE FROM %s WHERE 1=1 " % (table)
            conditions=self.add_conditions(conditions)
            if conditions is not None:
                sql+="AND %s " % conditions
        self.do(sql)

class PostgreSQL(DBA):

    def __init__(self,host='localhost',port=5432,user=None,passwd=None,db=None,coding='utf8',**kwargs):
        self.host=host
        self.port=int(port)
        self.user=user
        self.passwd=passwd
        self.db=db
        super(PostgreSQL,self).__init__(coding)
        
    def connect(self):
        self.connection=postgre.connect(host=self.host,port=self.port,user=self.user,passwd=self.passwd,db=self.db)
        self.set_coding(self.coding)
        
class DBI(object):
    drivers={
        'mysql':MySQL,
        'sqlite':SQLite,
        'postgre':PostgreSQL,
    }
    
    def __new__(self,driver='mysql',*args,**kwargs):
        cls=DBI.drivers.get(driver,None)
        if cls is None:
            raise Exception('No such dbi driver %s.' % driver)
        if 'dbc' in kwargs:
            db=kwargs.pop('dbc')
            if db in dbc:
                return cls(**dbc[db])
            else:
                raise Exception('No such db %s in dbc.' % db)
        return cls(*args,**kwargs)
    
def main():
    sql="select local_id,local_name,list_name from display_local "
    dbi=DBI(dbc='db_base_dict')
    for row in dbi.query(sql):
        print row
    print dbi.table_fields('display_local')
    print dbi.has_table('display_local')
    print dbi.has_db('dbwww58com_groupbuy_stats')
    print dbi.has_record('display_local',where=dict(list_name='bj'))
    for row in dbi.batch(sql,1000):
        print ','.join(str(field) for field in row)

if __name__ == '__main__':
    main()
